# 1. 导入需要的库
from gensim.utils import simple_preprocess
from gensim import corpora
from gensim.models import TfidfModel
import numpy as np
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report

# 2. 准备数
texts = [
    "I love machine learning",
    "Deep learning is a branch of machine learning",
    "Support vector machines are supervised learning models",
    "Natural language processing is fun",
    "I enjoy learning new things about AI"
]
labels = [0, 0, 1, 1, 0]  # 文本对应的分类标签

# 3. 文本预处理：分词
tokenized_texts = [simple_preprocess(doc) for doc in texts]

# 4. 建立词典
dictionary = corpora.Dictionary(tokenized_texts)

# 5. 文本转成BOW表示
corpus_bow = [dictionary.doc2bow(text) for text in tokenized_texts]

# 6. 建立TF-IDF模型
tfidf_model = TfidfModel(corpus_bow)
corpus_tfidf = [tfidf_model[doc] for doc in corpus_bow]

# 7. 稀疏向量转稠密矩阵
num_features = len(dictionary)

def sparse_to_dense(sparse_vec, num_features):
    dense_vec = np.zeros(num_features)
    for idx, value in sparse_vec:
        dense_vec[idx] = value
    return dense_vec

X = np.array([sparse_to_dense(doc, num_features) for doc in corpus_tfidf])
y = np.array(labels)

# 8. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 9. 建立SVM模型并训练
svm_clf = SVC(kernel='linear')
svm_clf.fit(X_train, y_train)

# 10. 预测并输出结果
y_pred = svm_clf.predict(X_test)
print(classification_report(y_test, y_pred))
